#include <chrono>
#include <eigen3/Eigen/Dense>
#include <iostream>
#include <vector>

// 原始的嵌套循环方法
Eigen::MatrixXd
distance_matrix_loop(const std::vector<Eigen::Vector3d> &detections,
                     const std::vector<Eigen::Vector3d> &trks) {
  Eigen::MatrixXd distance_matrix(detections.size(), trks.size());
  for (size_t i = 0; i < detections.size(); ++i) {
    for (size_t j = 0; j < trks.size(); ++j) {
      distance_matrix(i, j) = (detections[i] - trks[j]).norm();
    }
  }
  return distance_matrix;
}

// 矢量化方法
Eigen::MatrixXd
distance_matrix_vectorized(const std::vector<Eigen::Vector3d> &detections,
                           const std::vector<Eigen::Vector3d> &trks) {
  int num_detections = detections.size();
  int num_trks = trks.size();

  // 将 `detections` 和 `trks` 转换为 Eigen 矩阵
  Eigen::MatrixXd det_matrix(num_detections, 3);
  Eigen::MatrixXd trk_matrix(num_trks, 3);

  for (size_t i = 0; i < num_detections; ++i) {
    det_matrix.row(i) = detections[i].transpose();
  }
  for (size_t j = 0; j < num_trks; ++j) {
    trk_matrix.row(j) = trks[j].transpose();
  }

  // 矢量化计算距离矩阵
  Eigen::MatrixXd d_sq =
      det_matrix.rowwise().squaredNorm().replicate(1, num_trks);
  Eigen::MatrixXd t_sq =
      trk_matrix.rowwise().squaredNorm().transpose().replicate(num_detections,
                                                               1);
  Eigen::MatrixXd distance_matrix =
      (d_sq + t_sq - 2 * det_matrix * trk_matrix.transpose()).array().sqrt();

  return distance_matrix;
}

int main() {
  // 初始化一些随机数据用于测试
  int num_detections = 1000;
  int num_trks = 1000;

  std::vector<Eigen::Vector3d> detections(num_detections);
  std::vector<Eigen::Vector3d> trks(num_trks);

  for (auto &det : detections) {
    det = Eigen::Vector3d::Random();
  }
  for (auto &trk : trks) {
    trk = Eigen::Vector3d::Random();
  }

  // 测试嵌套循环方法的运行时间
  auto start = std::chrono::high_resolution_clock::now();
  Eigen::MatrixXd result_loop = distance_matrix_loop(detections, trks);
  auto end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> duration_loop = end - start;
  std::cout << "Loop method duration: " << duration_loop.count() << " seconds"
            << std::endl;

  // 测试矢量化方法的运行时间
  start = std::chrono::high_resolution_clock::now();
  Eigen::MatrixXd result_vectorized =
      distance_matrix_vectorized(detections, trks);
  end = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double> duration_vectorized = end - start;
  std::cout << "Vectorized method duration: " << duration_vectorized.count()
            << " seconds" << std::endl;

  // 验证结果一致性（可选）
  bool is_approx_equal = result_loop.isApprox(result_vectorized, 1e-6);
  std::cout << "Results are approximately equal: "
            << (is_approx_equal ? "Yes" : "No") << std::endl;

  return 0;
}
